from tjob.initialize.setup_envs import _SUBMIT_JOBS
import os.path as osp

from tjob.utils.io_utils import load_yaml, save_yaml
from tjob.utils.job_utils import JobInfo
import warnings
from tjob.handler import LocalHandler, PAIHandler

# return code

SUCCESS = 0
FAILED = 1
PART_SUCCESS = 2


class JobManger:
    """Job Manger. Record your submitted jobs."""

    def __init__(self) -> None:
        if not osp.exists(_SUBMIT_JOBS):
            self.meta_submit = {}
        else:
            self.meta_submit = load_yaml(_SUBMIT_JOBS)

    def add(self, data: JobInfo):
        name = data.name
        if name is None:
            _msg = f"[JobManger] Cannot add job's name is None, JobInfo: {data}"
            raise NameError(_msg)

        if name in self.meta_submit:
            new_name = name + "." + str(data.create)
            _msg = (
                f"[JobManger] name: {name} exits, "
                + f"add its create time in name: {new_name} "
            )
            warnings.warn(_msg)
            name = new_name

        self.meta_submit.update({name: data.dict()})
        save_yaml(_SUBMIT_JOBS, self.meta_submit)

    def _update(self, name, **kwargs):
        if name not in self.meta_submit:
            _msg = f"Update FAILED, {name} not in JobManger"
            raise ValueError(_msg)
        ori_data = self.meta_submit[name]
        ori_data.update(kwargs)
        save_yaml(_SUBMIT_JOBS, self.meta_submit)

    def remove(self, name, pop=True) -> bool:
        if name in self.names:
            job_info = self.meta_submit.get(name)
            if job_info["type"].strip().lower() == "pai":
                job_id = job_info["id"]
                # _msg = "Try remove job in PAI ... "
                ret = PAIHandler.remove(job_id)
                if ret == 0 and pop:
                    # sucesss remove in pai
                    self.meta_submit.pop(name)
                    save_yaml(_SUBMIT_JOBS, self.meta_submit)
                return ret == 0
            else:
                warnings.warn("local is already removed after Job finish")
                self.meta_submit.pop(name)
                save_yaml(_SUBMIT_JOBS, self.meta_submit)
                return True
        return False

    def pop(self, name):
        print("[POP] Only pop in JobManger, wont effect PAI")
        if name in self.names:
            job_info = self.meta_submit.get(name)
            print(f"[POP] {name}, info: {job_info}")
            self.meta_submit.pop(name)
            return True
        else:
            warnings.warn(f"{name} is not in JobManger")
            return False

    def clear(self):
        self.meta_submit = {}
        save_yaml(_SUBMIT_JOBS, self.meta_submit)

    def clone(self, name, update_stat=False) -> bool:
        """clone job"""
        submit_info = self.meta_submit.get(name, None)
        if submit_info is None:
            # warnings.warn(name + " is not in JobManger")
            return False
        submit_type = submit_info["type"]
        if submit_type.strip().lower() == "local":
            handler = LocalHandler(name, None)
            new_id = handler.info.id

            ret = handler.send(submit_info["command"])
            if ret == 0 and update_stat:
                # success will update job stats
                JobManger()._update(name, id=new_id)
            return ret == 0

        elif submit_type.strip().lower() == "pai":
            try:
                dlc_params = submit_info["dlc_params"]
            except Exception as e:
                print(e)
                raise ValueError("cannot get `dlc_params`")
            # print(dlc_params)
            ret = PAIHandler.params_submit(dlc_params)
            for line in ret.stdout:
                if "http" in str(line):
                    print("job url: " + str(line).strip().split("|")[-2].strip())
                    new_id = str(line).strip().split("|")[1].strip()
                    print("job id : " + new_id)
                    if update_stat:
                        JobManger()._update(name, id=new_id)
            ret.communicate()[0]
            return ret.returncode == 0
        else:
            raise NotImplementedError(submit_type)

    def stop(self, name):
        submit_info = self.meta_submit.get(name, None)
        if submit_info is None:
            # warnings.warn(name + " is not in JobManger")
            return False
        submit_type = submit_info["type"]
        id = submit_info["id"]
        if submit_type.strip().lower() == "local":
            ret = LocalHandler.stop(id)
            return ret == 0
        elif submit_type.strip().lower() == "pai":
            ret = PAIHandler.stop(id)
            return ret == 0
        else:
            raise NotImplementedError(submit_type)

    def restart(self, name):
        """stop and then clone"""
        ret = self.remove(name, pop=False)
        if ret:
            return self.clone(name, update_stat=True)
        else:
            return FAILED

    @property
    def names(self):
        return self.meta_submit.keys()

    @property
    def jobs(self):
        return self.meta_submit

    def fetch_id(self, name):
        if name in self.meta_submit:
            return self.meta_submit[name]["id"]
        else:
            warnings.warn(f"{name} is not in JobManger")

    def display(self):
        for key, value in self.meta_submit.items():
            print("-" * 30)
            print(f"{key}: {value}")
